import streamlit as st
import tempfile
import os
from langchain.memory import ConversationBufferMemory
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import PromptTemplate
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.agents import create_react_agent,AgentExecutor
from langchain_community.callbacks.streamlit import StreamlitCallbackHandler
from langchain_openai import ChatOpenAI
from langchain.tools.retriever import create_retriever_tool

root_dir = os.path.dirname(os.path.abspath(__file__))

#设置streamlit应用的页面标题和布局
st.set_page_config(page_title="文档问答", layout="wide")

#设置页面标题
st.title("文档问答")

#上传txt文件，允许上传多个文件
uploaded_files = st.sidebar.file_uploader("上传txt文件", accept_multiple_files=True)

if not uploaded_files:
    st.info("请上传txt文件")
    st.stop()

#实现检索器
@st.cache_resource(ttl="1h")
def configure_retriever(files):
    #读取上传的文件，并写入一个临时目录
    docs = []
    temp_dir = tempfile.TemporaryDirectory(dir=root_dir)
    for uploaded_file in files:
        file_path = os.path.join(temp_dir.name, uploaded_file.name)
        with open(file_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        loader = TextLoader(file_path, encoding="utf-8")
        docs.extend(loader.load())

    #对文档进行分块
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = text_splitter.split_documents(docs)

    #创建嵌入模型和向量数据库
    vectordb  = Chroma(
    collection_name="chroma_collection",
    embedding_function=OpenAIEmbeddings(model="text-embedding-3-large"),
    persist_directory=root_dir +"/vector_store/chroma_langchain_db")  # 指定持久化的目录，如果不存在该目录，则创建该目录；如果不设置该参数，就默认使用内存存储。

    vectordb.add_documents(splits)
    #创建检索器
    return vectordb.as_retriever()

#配置检索器
retriever = configure_retriever(uploaded_files)

#如果session_state中没有消息记录或者用户点击清空聊天记录按钮，则初始化消息记录
if "messages" not in st.session_state or st.sidebar.button("清空聊天记录"):
    st.session_state["messages"] = [{"role": "assistant", "content": "你好，我是文档问答助手。"}]

#加载历史聊天记录
for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

#创建用于文档检索的工具
tool = create_retriever_tool(
    retriever=retriever,
    name="文档检索",
    description="用于搜索文档的工具。输入一个查询，并返回一个包含相关文档的列表。"
)

tools = [tool]

#创建聊天消息历史记录
msgs = StreamlitChatMessageHistory()

#创建对话缓冲区内存
memory = ConversationBufferMemory(chat_memory=msgs, return_messages=True, memory_key="chat_history", output_key="output")

#指令模板
instructions = """
      您是一个设计用于查询文档、回答问题的代理。
      您可以使用文档检索工具，并基于检索内存来回答问题。
      您可能不查询文档就知道答案，但是您仍然应当查询文档来获取答案。
      如果您从文档中找不到任何信息用于回答问题，则只需返回”抱歉，这个问题我还不知道。“作为答案。
"""

#基础提示模板
base_prompt_template = """
       {instructions}
       
       TOOLS:
       you have access to the following tools:
       
       {tools}
       
       TO use a tool, please use the following format:
       Thought : Do I need to use a tool?
       Action: the action to take, should be one of [{tool_names}]
       Action Input: {input}
       Observation: the result of the action
       
       When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
       Thought: Do I need to use a tool?
       Response: [your response here]
       
       Begin!
       
       Previous conversation history:
       {chat_history}
       
       New input: {input}
       {agent_scratchpad}
      """

#创建基础提示模板
prompt_template = PromptTemplate.from_template(base_prompt_template)

#创建部分填充的提示模板
prompt = prompt_template.partial(
    instructions=instructions
)

#创建llm
llm = ChatOpenAI()

#创建代理
agent = create_react_agent(
    llm=llm,
    tools=tools,
    prompt=prompt
)

#创建代理执行器
agent_executor = AgentExecutor(agent=agent, tools=tools,memory=memory,verbose=True,handle_parsing_errors="没有从知识库检索到相似内容")

#创建聊天输入框
user_query = st.chat_input(placeholder="请开始提问吧！")

#如果用户输入了问题，则执行代理执行器
if user_query:
    st.session_state.messages.append({"role": "user", "content": user_query})
    st.chat_message("user").write(user_query)

    with st.chat_message("assistant"):
        # 创建Streamlist 回调处理器
        st_cd = StreamlitCallbackHandler(st.container())

        # agent执行过程日志回调显示在Sreamlit Container中
        config = {"callbacks": [st_cd]}
        # 执行代理执行器
        response = agent_executor.invoke({"input": user_query},config=config)
        # 添加助手消息到session_state
        st.session_state.messages.append({"role": "assistant", "content": response["output"]})
        # 显示助手响应
        st.write(response["output"])